/**
 * 
 */
package edu.umd.clip.lm.programs;

import edu.berkeley.nlp.util.*;

import java.io.*;
import java.net.*;
import java.util.*;
import java.util.concurrent.locks.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.text.*;

import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.*;
import edu.umd.clip.jobs.*;
import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.factors.Dictionary;
import edu.umd.clip.lm.storage.*;
import edu.umd.clip.lm.storage.remote.*;
import edu.umd.clip.lm.util.*;

import java.lang.management.*;

import com.sleepycat.je.DatabaseException;
import com.sleepycat.je.Environment;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ComputeNgramProbabilities {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-input", required = false, usage = "Evaluation data file (Default: stdin)")
		public String input;
        @Option(name = "-output", required = false, usage = "Output file (Default: stdout)")
		public String output;
        @Option(name = "-jobs", usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-forest", required = true, usage = "the decision tree forest")
		public String forest;        
        @Option(name = "-debug", usage = "decoder debug level (default: 0)")
        public int debug = 0;
        @Option(name = "-host", usage = "remote storage hostname or IP address (default: use local storage)")
        public String host = null;
        @Option(name = "-port", usage = "remove storage port (default: 2332)")
        public int port = 0;
        @Option(name = "-meminfo", usage = "track memory usage (default: false)")
        public boolean meminfo = false;
        @Option(name = "-coarse-threshold", usage = "coarse threshold, requires the new decoder (default: 1e-4)")
        public double coarse = -1;
        @Option(name = "-lc", usage = "lowercase the input")
        public boolean lower = false;
        @Option(name = "-uc", usage = "uppercase the input")
        public boolean upper = false;
        @Option(name = "-jerboa", required = false, usage = "use Jerboa storage (default: false)")
		public boolean useJerboa = false;        
        @Option(name = "-compact", required = false, usage = "use Compact storage (default: false)")
		public boolean useCompact = false;        
        @Option(name = "-bdb", required = false, usage = "use Berkeley DB storage (default: false)")
		public boolean useBDB = false;        
	}

	/**
	 * @param args
	 * @throws DatabaseException 
	 * @throws IOException 
	 * @throws InterruptedException 
	 * @throws ClassNotFoundException 
	 */
	public static void main(String[] args) throws DatabaseException, IOException, InterruptedException, ClassNotFoundException {
		final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
		
		MemoryUsage memuse;
		
        OptionParser optParser = new OptionParser(Options.class);
        final Options opts = (Options) optParser.parse(args, true);

        if (opts.upper && opts.lower) {
        	System.err.println("Lowercase AND uppercase? doesn't sound right.");
        	return;
        }
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: initial: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

        LMDecodingOptions lmOpts = new LMDecodingOptions();
        lmOpts.config = opts.config;
        lmOpts.debug = opts.debug;
        if (opts.jobs > 0) lmOpts.jobs = opts.jobs;
        if (opts.host != null) lmOpts.host = opts.host;
        if (opts.port > 0) lmOpts.port = opts.port;
        if (opts.coarse > 0) lmOpts.coarseThreshold = opts.coarse;
        if (opts.forest != null) lmOpts.forest = opts.forest;
        
        if (opts.host != null && opts.port > 0) lmOpts.storage = LMDecodingOptions.Storage.REMOTE;
        if (opts.useJerboa) lmOpts.storage = LMDecodingOptions.Storage.JERBOA;
        if (opts.useBDB) lmOpts.storage = LMDecodingOptions.Storage.BDB;
        if (opts.useCompact) lmOpts.storage = LMDecodingOptions.Storage.COMPACT;
        
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: initial: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

        LanguageModel.initDecoding(lmOpts);

        Experiment experiment = Experiment.getInstance();

		final InputParser parser = new PlainInputParser(experiment.getTupleDescription());
		/*	
		{
			long[] endTuples = new long[1];
			endTuples[0] = experiment.getTupleDescription().createEndTuple(); 
			parser.setEndTuples(endTuples);
		}
		*/
		final ForestModel forest = experiment.getForest(opts.forest);
		final NgramMarginalization marginalizer = new NgramMarginalization(forest);
		
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after loading LMs: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

		final MutableInteger extraWords = new MutableInteger(0);
		if (parser.getStartTuples() != null) {
			extraWords.set(extraWords.intValue() + parser.getStartTuples().length);
		}
		if (parser.getEndTuples() != null) {
			extraWords.set(extraWords.intValue() + parser.getEndTuples().length);
		}

		final DataSequencer<String> sequencer = new DataSequencer<String>(1);
		final byte wordFactorIdx = experiment.getTupleDescription().getMainFactorIndex();
		
		class EvalRunnable implements Runnable {
			Pair<Integer, String> pair;
			
			public EvalRunnable(int sentNo, String line) {
				pair = new Pair<Integer, String>(sentNo, line);
			}
			
			public void run() {
				String input = pair.getSecond();
				if (opts.lower) {
					input = input.toLowerCase();
				} else if (opts.upper) {
					input = input.toUpperCase();
				}
		    	long[] ngram = parser.parseSentence(input);
		    	
		    	// remove <s>'s in the beginning, they will be added implicitly
		    	for(byte i=0; i<ngram.length; ++i) {
		    		if (!Dictionary.isStart(FactorTuple.getValue(ngram[i], wordFactorIdx))) {
		    			if (i > 0) {
		    				ngram = Arrays.copyOfRange(ngram, i, ngram.length);
		    			}
		    			break;
		    		}
		    	}
		    	
		    	double logProb = marginalizer.computeNgramProbability(ngram);
		    	
				StringBuilder sb = new StringBuilder();
		    	sb.append(pair.getSecond());
		    	sb.append("\t");
		    	sb.append(Float.toString((float)logProb));

		    	pair.setSecond(sb.toString());
		    	
		    	//System.err.printf("done line %d, %s\n", pair.getFirst(), pair.getSecond());
		    	//System.err.flush();

		    	sequencer.putItem(pair);
				if (opts.meminfo) {
					MemoryUsage memuse = memoryBean.getHeapMemoryUsage();
					System.out.printf("MEMUSE: after evaluating sentence %s: %dM/%dM\n", pair.getFirst().toString(), memuse.getUsed()/1048576, memuse.getMax()/1048576);
					//System.out.println("STATS: " + eval.getStatistics());
				}

			}
		}
		
		Timer timer = null;
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after experiment.closeXML(): %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
			timer = new Timer(true);
			TimerTask task = new TimerTask() {
				public void run() {
					MemoryUsage memuse = memoryBean.getHeapMemoryUsage();
					DateFormat df = DateFormat.getDateTimeInstance(DateFormat.SHORT, DateFormat.LONG);
					System.out.printf("MEMUSE: %s: %dM/%dM\n", df.format(new Date()), memuse.getUsed()/1048576, memuse.getMax()/1048576);					
				}
			};
			timer.scheduleAtFixedRate(task, 0, 10000);
		}

		try {
			BufferedReader reader = IO.getReader(opts.input == null ? System.in : IO.getInputStream((opts.input)));
		    BufferedWriter writer = IO.getWriter(opts.output == null ? System.out : IO.getOutputStream((opts.output)));

		    JobManager manager = JobManager.getInstance();
		    JobGroup group = manager.createJobGroup("decoding");
		    
		    Thread outputPrinter = new Thread(new OutputPrinter(sequencer, writer));
		    outputPrinter.start();
		    
		    int lineNo = 0;
		    for(String line = reader.readLine(); line != null; line = reader.readLine()) {
		    	Job job = new Job(new EvalRunnable(++lineNo, line), "decoding sentence #"+Integer.toString(lineNo));
		    	job.setPriority(lineNo);
		    	manager.addJob(group, job);
		    }
		    sequencer.setLastItem(lineNo);
		    /*
		    try {
			    while(!sequencer.isDone()) {
			    	Pair<Integer, String> item = sequencer.getFirst();
			    	System.err.printf("printing line %d, %s\n", item.getFirst(), item.getSecond());
			    	System.err.flush();
			    	writer.append(item.getSecond());
			    	writer.newLine();
			    	writer.flush();
			    }
		    } catch(InterruptedException e) {
		    	e.printStackTrace();
		    }
		    */
		    //System.err.println("all jobs submitted, waiting for the outputPrinter to finish...");
		    //System.err.flush();
		    
		    outputPrinter.join();
		    
		    //System.err.println("outputPrinter done");
		    //System.err.flush();
		    
		    //writer.append(opts.lm + " stats: \n" + eval.getStatistics() + "\n");
		    writer.close();
		    
		    if (opts.debug > 0) {
		    	System.err.println(forest.getName() + " stats: \n" + forest.getDecoder().getStatistics() + "\n");
		    }
		} catch(IOException e) {
			e.printStackTrace();
		}
	}
	
	private static class OutputPrinter implements Runnable {
		private final DataSequencer<String> sequencer;
		private final BufferedWriter writer;
		/**
		 * @param sequencer
		 * @param writer
		 */
		public OutputPrinter(DataSequencer<String> sequencer,
				BufferedWriter writer) {
			this.sequencer = sequencer;
			this.writer = writer;
		}

		/* (non-Javadoc)
		 * @see java.lang.Runnable#run()
		 */
		@Override
		public void run() {
		    try {
			    while(!sequencer.isDone()) {
			    	Pair<Integer, String> item = sequencer.getFirst();
			    	if (item== null) return;
			    	//System.err.printf("printing line %d, %s\n", item.getFirst(), item.getSecond());
			    	//System.err.flush();
			    	writer.append(item.getSecond());
			    	writer.newLine();
			    	writer.flush();
			    }
		    } catch(InterruptedException e) {
		    	e.printStackTrace();
		    } catch (IOException e) {
				e.printStackTrace();
			}
			
		}
		
	}
}
